Skip to content

fix: remove seq_parallel + tp restriction in dtensor v2#1725

Closed
terrykong wants to merge 1 commit intomainfrom
test-fix-3
Closed

fix: remove seq_parallel + tp restriction in dtensor v2#1725
terrykong wants to merge 1 commit intomainfrom
test-fix-3

Conversation

@terrykong
Copy link
Copy Markdown
Collaborator

@terrykong terrykong commented Jan 6, 2026

What does this PR do ?

Add a one line overview of what this PR aims to accomplish.

Issues

List issues that this PR closes (syntax):

Usage

  • You can potentially add a usage example below
# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

  • Make sure you read and followed Contributor guidelines
  • Did you write any new necessary tests?
  • Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
  • Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

  • ...

@terrykong
fix: remove seq_parallel + tp restriction in dtensor v2
4b7fb62 
Signed-off-by: Terry Kong <terryk@nvidia.com>
@terrykong terrykong requested review from yfw and yuki-97 January 6, 2026 08:42
@terrykong terrykong requested review from a team as code owners January 6, 2026 08:42
@terrykong terrykong added CI:L1 Run doctests, unit tests, and functional tests r0.5.0 labels Jan 6, 2026
@coderabbitai
Copy link
Copy Markdown
Contributor

coderabbitai bot commented Jan 6, 2026

Warning

Rate limit exceeded

@terrykong has exceeded the limit for the number of commits that can be reviewed per hour. Please wait 2 minutes and 32 seconds before requesting another review.

⌛ How to resolve this issue?

After the wait time has elapsed, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

We recommend that you space out your commits to avoid hitting the rate limit.

🚦 How do rate limits work?

CodeRabbit enforces hourly rate limits for each developer per organization.

Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout.

Please see our FAQ for further information.

📥 Commits

Reviewing files that changed from the base of the PR and between 1720466 and 4b7fb62.

📒 Files selected for processing (1)
  • nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

@github-actions
Copy link
Copy Markdown

github-actions bot commented Jan 6, 2026

⚠️ File Consistency Check

Check based on commit: 4b7fb62 (PR #1725 from test-fix-3)

⚠️ DTensor Policy Worker Synchronization Warning

The file nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py was modified in this PR, but nemo_rl/models/policy/workers/dtensor_policy_worker.py was not updated.

Why this matters:
These files contain related DTensor policy worker implementations that should be kept synchronized to ensure consistency across different versions.

Action required:

  • Please review if the changes in nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py should also be applied to nemo_rl/models/policy/workers/dtensor_policy_worker.py
  • Update nemo_rl/models/policy/workers/dtensor_policy_worker.py if necessary to maintain consistency
  • If the files are intentionally different, please add a comment in the PR explaining why

Files to check:

  • Modified: nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py
  • Not modified: nemo_rl/models/policy/workers/dtensor_policy_worker.py

This check ensures that related file implementations remain synchronized across the codebase. If you believe this warning is incorrect or the files should intentionally differ, please add a comment explaining the reasoning.

@terrykong terrykong mentioned this pull request Jan 6, 2026
Draft
yuki-97
yuki-97 reviewed Jan 6, 2026
View reviewed changes
nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py
print(
"[WARNING]: sequence_parallel=True, but tp_size=1 which has no effect. Enable tp_size > 1 to use sequence parallelism."
)
elif sequence_parallel_enabled and tp_size > 1:
Copy link
Copy Markdown
Contributor

@yuki-97 yuki-97 Jan 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

shall we also remove this restriction in dtensor v1?

@terrykong terrykong marked this pull request as draft January 6, 2026 20:22
@terrykong
Copy link
Copy Markdown
Collaborator Author

terrykong commented Jan 6, 2026

moving to draft from what @RayenTian found

@terrykong
Copy link
Copy Markdown
Collaborator Author

terrykong commented Jan 8, 2026 •
edited
Loading

closing in favor of #1728

@terrykong terrykong closed this Jan 8, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@yuki-97 yuki-97 yuki-97 left review comments

@yfw yfw Awaiting requested review from yfw

Assignees

No one assigned

Labels

CI:L1 Run doctests, unit tests, and functional tests r0.5.0

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@terrykong @yuki-97